#  sniffer.rb
#
#  This file is part of myvnet.
#
#  Copyright (C) 2008  Filipe David Borba Manana,  fdmanana@gmail.com
#
#  myvnet is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  myvnet is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with myvnet.  If not, see <http://www.gnu.org/licenses/>.


require 'lib/packet/raw_socket'
require 'lib/exceptions'
require 'lib/protocols/ethernet'
require 'lib/protocols/arp'


module Packet
  
  
  # Packet sniffer.
  #
  class Sniffer

    attr_reader :interface
    attr_reader :addresses
  
    # === Synopsis
    #
    #   Packet::Sniffer.new(interface, addresses, packet_count = -1)
    #
    # === Args
    #
    # +interface+::
    #   Name of the network interface to use for packet sniffing.
    #   
    # +addresses+::
    #   Array of IP (IPv4) addresses the sniffer will capture packets
    #   destined to.
    #   
    # +network_proto+::
    #   Network protocol, as a lowercase string or symbol, to which the 
    #   given addresses belong to. As of now, only IPv4 is supported.
    #   
    # +packet_count+::
    #   Maximum number of packets to capture. Defaults to unlimited (-1).
    #
    # === Exceptions
    #
    # +SnifferError+::
    #   If the specified network protocol is not supported.
    #   
    # +RawSocketError+::
    #   If an error occurred while creating a Linux Raw Socket for
    #   packet sniffing.
    #
    # === Description
    #
    # Instantiates a network traffic sniffer.
    #
    def initialize(interface, addresses, network_proto, packet_count = -1)
      @net_proto = network_proto.to_sym
      
      unless [:ipv4].include? @net_proto
        raise SnifferError, "Unsupported network layer protocol"
      end      
      
      @addresses = addresses
      @interface = interface
      @packet_count = packet_count            
      
      # maps byte stream addresses to human friendly format addresses
      @str_addr = {}
      
      # holds number of packets received for each address
      @recv = {}
      
      @addresses.each do |a|
        bs = a
        
        case @net_proto
        when :ipv4
          bs = bs.scan(/\d{1,3}/)
          bs = bs.map { |b| b.to_i }
          bs = bs.pack 'C4'
        end
        
        @str_addr[bs] = a
        @recv[a] = 0
      end
      
      # promiscuous mode not needed since hardware address is not spoofed
      @socket = Packet::RawSocket.new @interface      
    end
    
    
    # === Synopsis
    #
    #   Packet::Sniffer#each_frame { |frame| ... }
    #
    # === Return
    #
    # No value.
    # 
    # === Exceptions
    # 
    # +RawSocketError+::
    #   If an error occurred while reading from the 
    #   underlying Linux Raw Socket.
    #   
    # +Interrupt+::
    #   If the packet capture was interrupted. For example,
    #   the process received signal INT (ctrl+c).
    #
    # === Description
    #
    # Yields captured network traffic, as instances of the class Ethernet::Packet, 
    # to the given block. 
    #
    def each_frame(&block)
      
      recv_total = 0
      
      begin
        while true do
          break if @packet_count >= 0 and recv_total >= @packet_count
          
          pass = false
          data = @socket.recv          
          frame = Ethernet::Packet.new data          
          
          case frame.ethertype
            
          when :arp
            arp_pkt = ARP::Packet.new(frame.payload)
            frame.payload = arp_pkt
            
            if ( arp_pkt.htype == :ethernet and arp_pkt.ptype == :ipv4 and 
                 @recv.has_key?(arp_pkt.tpa) )
              pass = true
              @recv[arp_pkt.tpa] += 1
            end
            
          when :ipv4
            # do not build IPv4 packet, an expensive operation if it is not
            # destined to one of the target IPv4 addresses
            dst = frame.payload[16,4]            
            friendly_addr = @str_addr[dst]
            
            if friendly_addr
              pass = true
              @recv[friendly_addr] += 1
            end
            
          end
          
          if pass
            recv_total += 1
            block.call frame
          end
          
        end
        
      rescue Interrupt
        raise
      rescue RawSocketError
        raise        
      ensure
        @socket.close
      end
      
    end
        
    alias each_packet each_frame
    alias each each_frame
    
    
    # === Synopsis
    #
    #   Packet::Sniffer#received
    #
    # === Return
    #
    # An hash whose keys are addresses, and values correspond to the number
    # of captured packets destinated to the respective addresses.
    #
    # === Description
    #
    # Returns the number of packets received by the sniffer for each address.  
    #
    def received
      @recv
    end
  
  end
  
end
